#!/usr/bin/python3
import os, sys, json
import numpy as np
import pandas as pd

#python3 extract_origin_data.py

# 项目地址
projects_path = '/mri_projects/ASD_Analysis/'
# 表文件地址
lists_path = os.path.join(projects_path, 'Lists')
# 预处理文件地址
csv_file = os.path.join(lists_path, 'Description/SubjectDescription.csv')

csv_data = pd.read_csv(csv_file)
# print(csv_data.columns)
# ASD_data = csv_data[csv_data.DX_GROUP == 1]
# GU_ASD_data = ASD_data[csv_data.SITE == 'GU']

# HC_data = csv_data[csv_data.DX_GROUP == 2]
# GU_HC_data = HC_data[csv_data.SITE == 'GU']

def print_mean_std(pd_data, clo_label, type, desc, round_num = 2):
    print("#" * 20)
    print("Statistic: [%s]" % desc)

    pt_str = "[%s] Count: %d, Range: %d~%d, Mean: %f, Std: %f, "

    if round_num == 2:
        pt_str += "Mean(Std): %.2f(%.2f)"
    if round_num == 3:
        pt_str += "Mean(Std): %.3f(%.3f)"

    print(pt_str % (type, len(pd_data[clo_label]), pd_data[clo_label].min(), pd_data[clo_label].max(), pd_data[clo_label].mean(), pd_data[clo_label].std(), round(pd_data[clo_label].mean(), round_num), round(pd_data[clo_label].std(), round_num)))

def statistic_age(ASD_data, HC_data):
    print_mean_std(ASD_data, "AGE_AT_SCAN", "ASD", "ASD_AGE", 2)
    print("-"*50)
    print_mean_std(HC_data, "AGE_AT_SCAN", "HC", "HC_AGE", 2)

def statistic_handedness(ASD_data, HC_data):
    ASD_hand_R = ASD_data[ASD_data.HANDEDNESS_CATEGORY == 1]
    if len(ASD_hand_R) > 0:
        print_mean_std(ASD_hand_R, "HANDEDNESS_CATEGORY", "ASD", "ASD_HAND_R", 2)

    ASD_hand_L = ASD_data[ASD_data.HANDEDNESS_CATEGORY == 2]
    if len(ASD_hand_L) > 0:
        print_mean_std(ASD_hand_L, "HANDEDNESS_CATEGORY", "ASD", "ASD_HAND_L", 2)

    ASD_hand_LR = ASD_data[ASD_data.HANDEDNESS_CATEGORY == 3]
    if len(ASD_hand_LR) > 0:
        print_mean_std(ASD_hand_LR, "HANDEDNESS_CATEGORY", "ASD", "ASD_HAND_LR", 2)

    print("-"*50)

    HC_hand_R = HC_data[HC_data.HANDEDNESS_CATEGORY == 1]
    if len(HC_hand_R) > 0:
        print_mean_std(HC_hand_R, "HANDEDNESS_CATEGORY", "HC", "HC_HAND_R", 2)

    HC_hand_L = HC_data[HC_data.HANDEDNESS_CATEGORY == 2]
    if len(HC_hand_L) > 0:
        print_mean_std(HC_hand_L, "HANDEDNESS_CATEGORY", "HC", "HC_HAND_L", 2)

    HC_hand_LR = HC_data[HC_data.HANDEDNESS_CATEGORY == 3]
    if len(HC_hand_LR) > 0:
        print_mean_std(HC_hand_LR, "HANDEDNESS_CATEGORY", "HC", "HC_HAND_LR", 2)

def statistic_fiq(ASD_data, HC_data):
    print_mean_std(ASD_data, "FIQ", "ASD", "ASD_FIQ", 2)
    print("-"*50)
    print_mean_std(HC_data, "FIQ", "HC", "HC_FIQ", 2)

def statistic_ADOS(ASD_data):
    head_list = ["ADOS_TOTAL", "ADOS_COMM", "ADOS_SOCIAL", "ADOS_STEREO_BEHAV", "ADOS_GOTHAM_SOCAFFECT", "ADOS_GOTHAM_RRB", "ADOS_GOTHAM_TOTAL", "ADOS_GOTHAM_SEVERITY"]

    for head_item in head_list:
        if head_item == 'ADOS_TOTAL':
            ADOS_Item = ASD_data[ASD_data.ADOS_TOTAL != -9999]
        elif head_item == 'ADOS_COMM':
            ADOS_Item = ASD_data[ASD_data.ADOS_COMM != -9999]
        elif head_item == 'ADOS_SOCIAL':
            ADOS_Item = ASD_data[ASD_data.ADOS_SOCIAL != -9999]
        elif head_item == 'ADOS_STEREO_BEHAV':
            ADOS_Item = ASD_data[ASD_data.ADOS_STEREO_BEHAV != -9999]
        elif head_item == 'ADOS_GOTHAM_SOCAFFECT':
            ADOS_Item = ASD_data[ASD_data.ADOS_GOTHAM_SOCAFFECT != -9999]
        elif head_item == 'ADOS_GOTHAM_RRB':
            ADOS_Item = ASD_data[ASD_data.ADOS_GOTHAM_RRB != -9999]
        elif head_item == 'ADOS_GOTHAM_TOTAL':
            ADOS_Item = ASD_data[ASD_data.ADOS_GOTHAM_TOTAL != -9999]
        elif head_item == 'ADOS_GOTHAM_SEVERITY':
            ADOS_Item = ASD_data[ASD_data.ADOS_GOTHAM_SEVERITY != -9999]

        print_mean_std(ADOS_Item, head_item, "ASD", "ASD_" + head_item, 2)

def statistic_FD(ASD_data, HC_data):
    print_mean_std(ASD_data, "meanFD", "ASD", "ASD_FD", 3)
    print("-"*50)
    print_mean_std(HC_data, "meanFD", "HC", "HC_FD", 3)

def get_variability_list(projects_path, extract_dir):
    lists_path = os.path.join(projects_path, 'Lists')

    asd_sub_file = os.path.join(lists_path, 'variability_list_ASD_' + extract_dir + '.txt')
    hc_sub_file = os.path.join(lists_path, 'variability_list_HC_' + extract_dir + '.txt')

    asd_list = load_list_file(asd_sub_file)
    hc_list = load_list_file(hc_sub_file)

    return asd_list, hc_list

def load_list_file(file_name):
    list_data = []

    fp = open(file_name, "r")
    for line in fp.readlines():
        list_data.append(line.strip())
    fp.close()

    return list_data

def statistic_center(csv_data, center_label = '', after_process = False):
    ASD_data = csv_data[csv_data.DX_GROUP == 1]
    HC_data = csv_data[csv_data.DX_GROUP == 2]

    if center_label != '':
        ASD_data = ASD_data[ASD_data.SITE == center_label]
        HC_data = HC_data[HC_data.SITE == center_label]

    if after_process:
        asd_list, hc_list = get_variability_list(projects_path, center_label)

        asd_sub_list = [sub_item.split('_')[-1] for sub_item in asd_list]
        hc_sub_list = [sub_item.split('_')[-1] for sub_item in hc_list]

        ASD_data = ASD_data[ASD_data["SUB_ID"].isin(asd_sub_list)]
        HC_data = HC_data[HC_data["SUB_ID"].isin(hc_sub_list)]

    print("*"*20 + " AGE " + "*"*20)
    statistic_age(ASD_data, HC_data)
    print("*"*50 + "\n")

    print("*"*20 + " HAND " + "*"*20)
    statistic_handedness(ASD_data, HC_data)
    print("*"*50 + "\n")

    print("*"*20 + " FIQ " + "*"*20)
    statistic_fiq(ASD_data, HC_data)
    print("*"*50 + "\n")

    print("*"*20 + " ADOS " + "*"*20)
    statistic_ADOS(ASD_data)
    print("*"*50 + "\n")

    print("*"*20 + " FD " + "*"*20)
    statistic_FD(ASD_data, HC_data)
    print("*"*50 + "\n")

statistic_center(csv_data, "KKI", True)
